package com.zzw.redislimter;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import reactor.core.publisher.Mono;

import java.util.Arrays;
import java.util.List;

/**
 * @author Administrator
 */
public class RRedisRateLimiter implements RLimter {

    private static final Logger logger = LoggerFactory.getLogger(RedisRateLimiter.class);
    private final StringRedisTemplate redisTemplate;
    private final RedisScript<List<Long>> script;

    public RRedisRateLimiter(StringRedisTemplate redisTemplate, RedisScript<List<Long>> script) {
        this.redisTemplate = redisTemplate;
        this.script = script;
    }

    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    @Override
    public Mono<Response> isAllowed(String key, String rate, String crust) {
        List<String> keys = getKeys(key);
        List<Long> exec;
        try {
            exec = this.redisTemplate.execute(this.script, keys, rate, crust, "", "1");
        } catch (Throwable throwable) {
            logger.error("Error calling rate limiter lua", throwable);
            exec = Arrays.asList(1L, -1L);
        }
        assert exec != null;
        boolean allowed = exec.get(0) == 1L;
        return Mono.just(new Response(allowed));
    }
}
